import torch
import torch.nn as nn

try:
    from .yolov1_basic import Conv
except:
    from yolov1_basic import Conv


class DecoupledHead(nn.Module):
    def __init__(self, cfg, in_dim, out_dim, num_classes=80):
        super(DecoupledHead, self).__init__()
        # =========类别检测头特征==============
        cls_out_dim = max(out_dim, num_classes)
        num_cls_head = cfg["num_cls_head"]
        cls_feas = []
        for i in range(num_cls_head):
            if i == 0:
                cls_feas.append(
                    Conv(
                        in_dim,
                        cls_out_dim,
                        3,
                        1,
                        1,
                        1,
                        cfg["head_act"],
                        cfg["head_norm"],
                    )
                )
            else:
                cls_feas.append(
                    Conv(
                        cls_out_dim,
                        cls_out_dim,
                        3,
                        1,
                        1,
                        1,
                        cfg["head_act"],
                        cfg["head_norm"],
                    )
                )
        # =========回归检测头特征==============
        reg_out_dim = max(out_dim, 64)
        num_reg_head = cfg["num_reg_head"]
        reg_feas = []
        for i in range(num_reg_head):
            if i == 0:
                reg_feas.append(
                    Conv(
                        in_dim,
                        reg_out_dim,
                        3,
                        1,
                        1,
                        1,
                        cfg["head_act"],
                        cfg["head_norm"],
                    )
                )
            else:
                reg_feas.append(
                    Conv(
                        reg_out_dim,
                        reg_out_dim,
                        3,
                        1,
                        1,
                        1,
                        cfg["head_act"],
                        cfg["head_norm"],
                    )
                )
        self.cls_feas = nn.Sequential(*cls_feas)
        self.reg_feas = nn.Sequential(*reg_feas)

    def forward(self, x):
        cls_fea = self.cls_feas(x)
        reg_fea = self.reg_feas(x)
        return cls_fea, reg_fea


def build_head(cfg, in_dim, out_dim, num_classes):
    model = cfg["head"]
    print("=======head=========")
    print("model head:{}".format(model))
    if model == "decoupled_head":
        return DecoupledHead(cfg, in_dim, out_dim, num_classes)
    else:
        raise NotImplementedError("head {} not implemented".format(model))
